library(data.table)
library(mefa)
library(doParallel)
library(emayili)
library(magrittr)



source('helper_functions.r')


runSensitivity <- function(config) {
  
  # Load the estimation
    print(paste0('Loading estimation result ',config))
    estimation <- loadEstimation(config) 
    
    
    dir.create(paste0('estimation_moment_sensitivities_jun8/',estimation$pars$filename,'_sensitivities/'))
    
    
    moments <- pars$moment_weights
    runMomentSensitivity(estimation,moments,'with_jacobian')
}


runMomentSensitivity <- function(estimation.results,momentWeights,filename) {
  
  markets.data <- estimation.results$market.data
  d.in <- fread('../data/final/estimation_data_sep9.csv')
  markets.data <- markets.data[,names(d.in),with=F]
  
  
  pars <- estimation.results$pars
  pars$CORES <- 10
  
  # Create instrumented interest rate
  rate.model <- lm(RATE ~ JUMBO * (D60.L1 ) + as.factor(paste0(YEAR,CBSA))   , data = markets.data,weights=markets.data$MARKET_SIZE)
  markets.data[,RATE_INST := predict(rate.model,markets.data)]
  
  # Set initial guess and raw draws. Initial guess is just the solution of the model
  new_mx <- param_vec_to_matrix(x = estimation.results$solution)
  pars$A <- new_mx$A
  pars$B <- new_mx$B
  pars$S <- new_mx$S
  raw.draws <- createRawDraws(pars)
  delta.init <<- NULL
  
  # Pre-loop
  # Calculate weights (if needed)
  # Initialize delta and drop some unsolveable markets.
  
  temp.pars <- pars
  temp.pars$BLP_TOL = 1e-10
  temp.pars$BLP_ATTEMPTS = 1000
  contraction <- solveAllMarkets(markets.data,raw.draws,temp.pars,delta.init,cores = temp.pars$CORES)
  contraction <- contraction[converged == 1]
  markets.data <- markets.data[MARKET_ID %in% contraction$MARKET_ID]
  delta.init <<- contraction[,c('MARKET_ID','j','delta')]
  
  moment_weights <- momentWeights
  
  # Calculate GMM objective funcion
  gmm_objective <- function(x,diag = F) {
    
    
    print(x)
    
    # 1.  Deal parameters
    new_mx <- param_vec_to_matrix(x)
    pars$A <- new_mx$A
    pars$B <- new_mx$B
    pars$S <- new_mx$S
    
    # 2. Recover deltas
    contraction <- solveAllMarkets(markets.data,raw.draws,pars,delta.init,cores = pars$CORES)
    
    # 3. Update the initial guesses for the good ones.
    next.init <- contraction
    next.init[converged == 0,delta := mean(next.init[converged == 1]$delta,na.rm=T)]
    delta.init <<- next.init[,c('MARKET_ID','j','delta')]
    
    # 4. Merge the deltas to the main marketdata
    market.data.merged <- merge(markets.data,contraction,by=c('MARKET_ID','j'))
    
    # 5. Regress out the remaining parameters
    if(length(pars$EST_YEARS) > 1) {
      linear <- lm(delta ~ RATE_INST + JUMBO + as.factor(YEAR) * PURPOSE * LENDER,weights = market.data.merged$converged,data = market.data.merged)
    } else {
      linear <- lm(delta ~ RATE_INST + JUMBO + PURPOSE * LENDER,weights = market.data.merged$converged,data = market.data.merged)
    }
    
    # 6. Calculate moments, apply wieghts, and calculate objective function
    moments <- calculateMoments(market.data.merged,linear)
    moments[,value := model - data]
    moments <- merge(moments,moment_weights,by='moment')
    moments[is.na(value),value := 999999]
    moments[is.infinite(value),value := 999999]
    
    val     <- sum(moments$value^2 * moments$weight)
    if(is.na(val)) {
      val <- 1e12
    }
    
    
    
    # 7. Report diagnostics if we want
    if(pars$diag.show_linear) {
      print(linear)
    }
    if(pars$diag.show_moments) {
      moments[,contribution := value^2 * weight]
      print(moments)  
    }
    if(pars$diag.show_convergence) {
      print(paste0(sum(contraction$converged)/4,'/',nrow(contraction)/4,' markets converged. Mean iterations given convergence = ',mean(contraction[converged == 1]$attempts)))
    }
    
    
    # 8. Report actual objective function
    print(val)
    
    if(!diag) {
      return(val)
    } else {
    
    
      return(list(x = x, market.data.merged = market.data.merged, linear.model = linear, moments = moments, val = val))
    }  
    
  }
  
  # Run the optimization
  scale = abs(estimation.results$solution) # 1e-1 * c(10,1,.1,.1,.1,.1,.1,.1,.1,.1,.1,.1,.1,.1)
  l = estimation.results$solution - .1 * abs(estimation.results$solution)
  u = estimation.results$solution + .1 * abs(estimation.results$solution)
  # Var used t obe dividing by 10, not 2.5
  
  
  # Calculate the jacobian of g(theta), e.g., the stacked moments, as you change the value.....
  for(xx in 1:length(estimation.results$solution)) {
    #xInput = result$par
    print(xx)
    xInput <- estimation.results$solution
    xInput[xx] <- xInput[xx] + .01
    dd <- gmm_objective(xInput,diag=T)
    JJ_temp <- dd$moments
    JJ_temp$param_idx <- xx
    if(xx == 1) {
      JJ <- JJ_temp
    } else {
      JJ <- rbind(JJ,JJ_temp)
    }
  }
  
  
  
  
  save(list = c('JJ','estimation.results'),file = paste0('estimation_moment_sensitivities_jun8/',estimation.results$pars$filename,'_sensitivities/',filename,'.rsave'))
  
  
}

calculateSensitivity <- function() {
  
  # load estimation
  load('estimation_moment_sensitivities_jun8/estimation_24_sensitivities/with_jacobian.rsave')
  
  
  dx = .01
  
  # Merge on basic stuff in order to calculate the jacobian
  JJ <- merge(JJ,estimation.results$moments[,c('moment','value')],by=c('moment'),suffixes = c('','.baseline'))
  JJ[,dMdX := (value - value.baseline)/dx]
  JJ <- JJ[order(moment,param_idx)]
  G = matrix(data = JJ$dMdX,nrow = length(unique(JJ$moment)),ncol = length(unique(JJ$param_idx)),byrow = T)
  
  JJ_moments <- JJ[!duplicated(moment),c('moment','weight')]
  JJ_moments <- JJ_moments[order(moment)]
  W = diag(JJ_moments$weight)
  
  Lambda = -1 * solve(t(G) %*% W %*% G) %*% t(G) %*% W 
  
  
  Lambda = -1 * solve(t(G) %*%  %*% G) %*% t(G) 
  
}



execute <- function() {
  ## Execute it.
  # Load config parameters
  config <- 'estimation_config'
  
 
  runSensitivity(config)
 
}



